#include "Splitter.h"

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/visitors.hpp>
#include <boost/graph/depth_first_search.hpp>
#include "Log.h"
#include "PS_Log.h"
#include <iostream>
#include <sstream>
#include <string>


Splitter::Splitter(Skinner *skinner) {
    this->skinner = skinner;
}

//
// Build FNF, using the given set of extreme objects
//
fnf_graph_t *Splitter::
build_farthest_neighbor_graph(const LMO_Span& span,
                              const LMO_Span& extreme_set,
                              Antipole_Map& antipoles) {
    LOG_ENTER(SPLIT,"*Splitter:: build_farthest_neighbor_graph");
    LOG(SPLIT) << "here is the cluster:";
    span.write(Log::output);
    Log::output.flush();
    LOG(SPLIT) << "here is the extreme set:";
    extreme_set.write(Log::output);
    Log::output.flush();
    LOG(SPLIT) << "here are the antipoles:";
    for (Antipole_Map::iterator iter = antipoles.begin(); iter != antipoles.end(); iter++) {
        int tail_ID = (*iter).first;
        ID_Dist head = (*iter).second;
        LOG(SPLIT) << tail_ID << " -> " << head.ID << " (dist=" << head.dist << ")";
    }
    if (antipoles.begin() == antipoles.end())
        LOG(SPLIT) << "  (NO antipoles)";

    Log::output.flush();

    fnf_graph_t *fnf = new fnf_graph_t();

    // Step 1: add the vertices
    LMO_Iterator iter = span.start;
    int n = span.n;
    for (int i=0; i<n; i++) {
        Labeled_Metric_Object *item = *iter++;

        LOG(SPLIT) << "  got items[" << i << "]=item";

        FNF_Vertex_Info vi(item->x, item->y, 0);

        LOG(SPLIT) << "  which has (x y) = ("
                     << item->x << " "
                     << item->y << ")";

        fnf->add_vertex(vi, item->ID);
    }

    // 2: If the antipoles aren't known, compute them.
    if (antipoles.empty()) {
        iter = span.start;

        LOG(SPLIT) << "  looking for antipoles";

        for (int i=0; i<n; i++) {
            double max_dist = -1;
            Labeled_Metric_Object *item = *iter++;
            int src_ID = item->ID;

            LOG(SPLIT) << "  item(i=" << i << ")=" << item;

            LMO_Iterator jter = extreme_set.start;
            int n_extreme     = extreme_set.n;
            int antipole_ID;

            for (int j=0; j<n_extreme; j++) {
                if (iter == jter) continue; // don't compare obj to itself

                Labeled_Metric_Object *jtem = *jter++;

                LOG(SPLIT) << "    jtem(j=" << j << ")=" << jtem;

                double dist = item->distance(jtem);
                if (dist > max_dist) {
                    max_dist = dist;
                    antipole_ID = jtem->ID;
                }
            }

            antipoles.insert(std::make_pair(src_ID, ID_Dist(antipole_ID, max_dist)));
        }
    }

    // 3: Finally, add the edges of the graph.
    LOG(SPLIT) << "  adding edges";

    iter = span.start;
    for (int i=0; i<n; i++) {
        Labeled_Metric_Object *item = *iter++;
        int src_ID = item->ID;

        LOG(SPLIT) << "getting antipole";

        ID_Dist& atem = antipoles.at(src_ID);
        int    antipole_ID  = atem.ID;
        double distance     = atem.dist;

        LOG(SPLIT) << "antipole ID=" << antipole_ID;

        FNF_Edge_Info ei(distance);

        LOG(SPLIT) << "adding edge " << src_ID << " > " << antipole_ID << " " << ei;

        fnf->add_edge(src_ID, antipole_ID, ei);
    }

    LOG(SPLIT) << "built fnf, which is:";
    fnf->print(Log::output);
    Log::output.flush();

    LOG_EXIT();

    return fnf;
}

// Class used by boost DFS called in color_fnf
class fnf_dfs_colorizer : public boost::default_dfs_visitor {
public:
    // This is invoked whenever DFS probes an edge.
    template <typename E, typename G>
    void examine_edge(E e, G& g) {


        LOG_ENTER(COLOR,"examine_edge");
        LOG(COLOR) << "  e=" << e;

        visit_edge((fnf_edge_t)e); //, (fnf_graph_t::graph_t&) g);

        LOG_EXIT();
    }

    void visit_edge(fnf_edge_t e) { //, fnf_graph_t::graph_t& g) {

        LOG_ENTER(COLOR, "visit_edge");
        LOG(COLOR) << "  e=" << e;
        // LOG(COLOR) << " let's dump the graph:";
        // LOG(COLOR) << " which has " << fnf->get_vertex_count() << " vertices"
        //                   << " and " << fnf->get_edge_count() << " edges";
        // LOG(COLOR) << *fnf;

        fnf_vertex_t u = fnf->get_source(e);
        fnf_vertex_t v = fnf->get_target(e);

        LOG(COLOR) << " u=" << u << " v=" << v;

        int u_name = fnf->get_vertex_name(u);
        int v_name = fnf->get_vertex_name(v);

        LOG(COLOR) << " u_name=" << u_name << " v_name=" << v_name;

        if (*pass == 1) {
            // There are two kinds of "root" vertices.
            // 1. those that occur in pairs, such that u -> v -> w is a cycle,
            // with w == u.  Then just color one of u,v white, and
            // the other black (for predictability, color the smaller one
            // white).
            // 2. Isolated roots, such that u -> v -> w is NOT a cycle,
            // ie w != u.  I will defer coloring these until
            // subsequent DFS passes.
            fnf_adjacency_vertex_range_t r = fnf->get_adjacent_vertices(v);
            fnf_adjacency_iter_t iter     = r.first;
            fnf_adjacency_iter_t iter_end = r.second;

            LOG(COLOR) << (iter == iter_end ? "v has no neighbors"
                           : "v has neighbors");

            fnf_vertex_t w = *iter;

            LOG(COLOR) << " w=" << w;

            if (w == u) {
                LOG(COLOR) << " u=" << u << " is root";

                // u or v is root, depending on order of tree construction
                if (u < v) {
                    LOG(COLOR) << " u gets WHITE";

                    fnf->properties(u).color = Color::WHITE;
                    fnf->properties(v).color = Color::BLACK;
                    (*n_colored) += 2;
                }
                else {
                    LOG(COLOR) << " v gets WHITE";

                    fnf->properties(u).color = Color::BLACK;
                    fnf->properties(v).color = Color::WHITE;
                }
            }
            else {
                LOG(COLOR) << " u=" << u << " is NON-root";

                // u's a non-root.  Mark it as unknown until the second pass.
                fnf->properties(u).color = Color::UNKNOWN;
                // don't increment n_colored.
            }
        }

        else {
            // non-first pass: give vertex the color opposite to its neighbor
            if (fnf->properties(u).color == Color::UNKNOWN &&
                fnf->properties(v).color != Color::UNKNOWN) {

                fnf->properties(u).color = 1 - fnf->properties(v).color;
                (*n_colored)++;
            }
        }

        LOG(COLOR) << "n_colored=" << *n_colored;
        LOG_EXIT();
    }

    int *pass;
    int *n_colored;
    fnf_graph_t *fnf;
};

// Hmmm... I want to use boost::depth_first_search, to visit all
// the vertices of the fnf, and find a 2-coloring of the fnf.
// The fnf is, after all, a forest (I've proved this).  So, we'll
// do a multi-pass algorithm:
// 1st pass: run dfs, and find all the "root vertices".  These come
//           in pairs which point to each other, and will be initialized
//           in black and white to seed the second pass.
// later pass: color vertices greedily, as they are visited, to be opposite
//           of their ancestor in the dfs tree

void Splitter::color_fnf(fnf_graph_t *fnf) {
    LOG_ENTER(SPLIT, "Splitter::color_fnf");
    LOG(SPLIT) << " coming in, fnf has " << fnf->get_vertex_count() << " vertices"
                 << " and " << fnf->get_edge_count() << " edges";

    fnf_dfs_colorizer colorizer;

    LOG(SPLIT) << "black is " << Color::BLACK
               << ", white is " << Color::WHITE
               << ", and unknown is " << Color::UNKNOWN;

    LOG(SPLIT) << " first DFS pass";

    int pass = 1;
    int n_colored = 0;

    colorizer.pass      = &pass;
    colorizer.n_colored = &n_colored;
    colorizer.fnf = fnf;
//    const fnf_graph_t::graph_t *container = fnf->get_graph();
    depth_first_search(fnf->get_graph(), visitor(colorizer));

    LOG(SPLIT) << "after 1st pass, here's the graph with twin roots colored:";
    LOG(SPLIT) << "\n" << *fnf;

    if (n_colored == 0) {
        LOG(ERROR) << "FATAL: First color pass found no twin roots!";
        LOG(ERROR) << "\n" << *fnf;
        exit(0);

        // If this turns out to be possible, need a fallback here:
        // pick a root, make it white, go from there.
    }

    while (n_colored < fnf->get_vertex_count()) {
        pass++;
        LOG(SPLIT) << " starting DFS pass #" << pass;
        LOG(SPLIT) << " before coloring, " << n_colored
                   << " done, should have done " << fnf->get_vertex_count();

        depth_first_search(fnf->get_graph(), visitor(colorizer));

        LOG(SPLIT) << "after pass #" << pass
                   << ", have " << n_colored
                   << " done, should have done " << fnf->get_vertex_count();
        LOG(SPLIT) << "Here's the graph fully colored:";
        LOG(SPLIT) << "\n" << *fnf;
    }
    LOG_EXIT();
}

int Splitter::partition(fnf_graph_t *fnf,  LMO_Span& span,
                        LMO_Iterator& start_of_white_set) {
    LOG_ENTER(SPLIT, "Splitter::partition");

    int i;
    LMO_Iterator iter;
    LMO_List black_set;  // used for partitioning
    LMO_List white_set;  // also for partitioning

    iter = span.start;
    int n = span.n;
    int n_white = 0;
    int n_black = 0;
    for (i=0; i<n; i++) {
        Labeled_Metric_Object *pt = *iter++;
        fnf_vertex_t v = fnf->get_vertex(pt->ID);
        FNF_Vertex_Info vi = fnf->properties(v);

        if (vi.color == Color::WHITE) {
            white_set.push_back(pt);
            n_white++;
        }
        else {
            black_set.push_back(pt);
            n_black++;
        }
    }

    //
    // Did the partition.  Now copy back into original
    // span of sites.
    //
    iter = span.start;

    LMO_Iterator biter = black_set.begin();
    for (i = 0; i < n_black; i++) {
        *iter++ = *biter++;
    }

    start_of_white_set = iter;
    LMO_Iterator witer = white_set.begin();
    for (i = 0; i < n_white; i++) {
        *iter++ = *witer++;
    }
    return n_black;

    LOG_EXIT();
}

void Splitter::split(Cluster *parent, Cluster*& A, Cluster*& B) {
    LOG_ENTER(SPLIT,"Splitter::split");

    LMO_Span span = parent->span;

    skinner->separate(parent);
    fnf_graph_t *fnf =
        build_farthest_neighbor_graph(span,
                                      parent->extreme_set.get_whole_span(),
                                      parent->antipoles);

    LOG(SPLIT) << "Before coloring, this is the parent cluster:";
    parent->write(Log::output);

    color_fnf(fnf);

    PS_Log::print(parent->span, false);
    std::ostringstream label;
    label << "CLUSTER ID " << parent->ID;
    PS_Log::show_title(label.str());
    PS_Log::ps->end_page();

    PS_Log::show_title("COLORED GRAPH");
    PS_Log::print_fnf(*fnf);

    LMO_Iterator pivot;
    int black_set_size = partition(fnf, span, pivot);

    LMO_Iterator black_end = pivot;
    black_end--;

    A = new Cluster(span.start, black_end, black_set_size,          0);
    B = new Cluster(pivot,      span.end,  span.n - black_set_size, 1);

    delete fnf;

    LOG_EXIT();
}
